import cv2
import glob
import os.path as osp
import sys

def showImg(img,wati_time=0,win_name="r"):
    cv2.namedWindow("r",cv2.WINDOW_NORMAL)
    cv2.imshow(win_name,img)
    key = cv2.waitKey(wati_time) & 0xFF
    if key == ord('q'):
        sys.exit()
    return key

def getKD(path):

    fs  = cv2.FileStorage(path, cv2.FileStorage_READ)
    K = fs.getNode('K').mat()
    D = fs.getNode('D').mat()
    return K,D

data_root = "/cv2/download_feishu/occ_data/tr_org"
img_list  = glob.glob(osp.join(data_root,"*jpg"))
img_list.sort()

tr_KD_p = "/home/westwell/welldriver/src/perception/utilitis/sync_image_ldiar_save/scripts/top_right/top_rightKD.yaml"

K,D = getKD(tr_KD_p)

for i,img_p in enumerate(img_list):
    img = cv2.imread(img_p)

    undistorted_img = cv2.undistort(img, K, D)
    undis_img_p = img_p.replace("tr_org","tr_undis")
    print(i,undis_img_p)
    cv2.imwrite(undis_img_p+"undis.jpg",undistorted_img)

    # showImg(img)
    # showImg(undistorted_img)

    